//
// ppu.c
//
// CSE 420
// Project 4 : Part 3
// Authors: Colin Hofman, Alan Tu, Adam Fairfield
//

// System includes
#include <stdio.h>
#include <stddef.h>
#include <stdint.h>
#include <ctype.h>
#include <stdlib.h>
#include <errno.h>
#include <string.h>
#include <libspe2.h>
#include <cbe_mfc.h>
#include <pthread.h>
#include <unistd.h>
#include <math.h>

// Local includes
#include "common.h"

// Data structure for running SPE thread
typedef struct ppu_pthread_data
{
	spe_context_ptr_t   speid;
	pthread_t           pthread;
	uint64_t						*argp;
} ppu_pthread_data_t;

// Data structure for keeping track of block indices in C
// running on a given SPE thread
typedef struct thread_indices
{
   int i_block;
   int j_block;
} thread_indices_t;

// Global variables
int* matrixA_InPPE                  QWORD_ALIGNED;    // Entire matrix a
int* matrixB_InPPE                  QWORD_ALIGNED;    // Entire matrix b
int* matrixC_InPPE                  QWORD_ALIGNED;    // Entire result matrix c
int a[NUM_SPUS][BLOCK_SIZE_SQUARED] QWORD_ALIGNED;    // Single block of matrix a to feed to spu
int b[NUM_SPUS][BLOCK_SIZE_SQUARED] QWORD_ALIGNED;    // Single block of matrix b to feed to spu
int c[NUM_SPUS][BLOCK_SIZE_SQUARED] QWORD_ALIGNED;    // Single block of matrix c to feed to spu
extern spe_program_handle_t spu;                      // Reference to spu program thread
CONTROL_BLOCK cb[NUM_SPUS]          QWORD_ALIGNED;		// Control block structure to pass to spu
ppu_pthread_data_t data[NUM_SPUS];							// PPU program thread data
thread_indices_t spu_indices[NUM_SPUS];               // Matrix block indices per thread running

// Forward function declarations
void *ppu_pthread_function(void *arg);
void initMatrices();
void printMatrix(char *name, int *matrix);

//
// Main ppu function.
//
// @param   int   argc     Argument count
//          char* argv[]   Argument list
// @return  int            Return status
//
int main(int argc, char *argv[])
{
	int rc[NUM_SPUS];

	// Initialize A and B with random values
	initMatrices();
   
  // Print matrices A and B
	//printMatrix("A", matrixA_InPPE);
	//printMatrix("B", matrixB_InPPE);

#ifdef TIMING
	long unsigned int start_t, end_t;
	float total_t = 0;
#endif

   int thread = 0;
   int nOverBlockSize = (N / BLOCK_SIZE);
   int remainingThreads = nOverBlockSize * nOverBlockSize * nOverBlockSize;
   int numThreads = (remainingThreads >= NUM_SPUS) ? NUM_SPUS : remainingThreads;

   // Multiply matrices using block multiplication
   for (int i_block = 0; i_block < N; i_block+=BLOCK_SIZE)
   {
      for (int j_block = 0; j_block < N; j_block+=BLOCK_SIZE)
      {
			for (int k_block = 0; k_block < N; k_block+=BLOCK_SIZE)
			{
            if (thread < numThreads)
            {
               //
               // Assemble a new thread
               //

               // Pull out individual block chunks
					for (int i = 0; i < BLOCK_SIZE; ++i)
					{
					   for (int j = 0; j < BLOCK_SIZE; ++j)
					   {
						   for (int k = 0; k < BLOCK_SIZE; ++k)
						   {
							   a[thread][i * BLOCK_SIZE + k] = matrixA_InPPE[(i_block + i) * N + (k_block + k)];
							   b[thread][k * BLOCK_SIZE + j] = matrixB_InPPE[(k_block + k) * N + (j_block + j)];
							   c[thread][i * BLOCK_SIZE + j] = matrixC_InPPE[(i_block + i) * N + (j_block + j)];
						   }
					   }
					}

               // Pack block data into control block
					cb[thread].blockA = &a[thread];
					cb[thread].blockB = &b[thread];
					cb[thread].SPU_result = &c[thread];
							
					// Create SPE context and load SPE program into the SPE context
					if ((data[thread].speid = spe_context_create(0, NULL)) == NULL)
					{
						fprintf(stderr, "Failed spe_context_create(errno=%d strerror=%s)\n", errno, strerror(errno));
						exit(3);
					}
			
					// Load program
					if ((rc[thread] = spe_program_load(data[thread].speid, &spu)) != 0)
					{
						fprintf(stderr, "Failed spe_program_load(errno=%d strerror=%s)\n", errno, strerror(errno));
						exit(3);
					}

					// Pass Control Block address to SPU
					data[thread].argp = (uint64_t*)&cb[thread];
               spu_indices[thread].i_block = i_block;
               spu_indices[thread].j_block = j_block;

               thread++;
            }

            if (thread == numThreads)
            {
               
#ifdef TIMING
	            start_t = __mftb();
#endif
               //
               // Run all threads and wait for them to complete
               //
   
					// Create SPE pthreads
               for (int t = 0; t < numThreads; t++)
					{
						if ((rc[t] = pthread_create(&data[t].pthread, NULL, &ppu_pthread_function, &data[t])) != 0)
						{
							fprintf(stderr, "Failed pthread_create(errno=%d strerror=%s)\n", errno, strerror(errno));
							exit(3);
						}
					}
               
					// Wait for SPE to complete
               for (int t = 0; t < numThreads; t++)
					{
						if ((rc[t] = pthread_join(data[t].pthread, NULL)) != 0)
						{
							fprintf(stderr, "Failed pthread_join(rc=%d, errno=%d strerror=%s)\n", rc[t], errno, strerror(errno));
							exit(1);
						}
					}
						
               // Store block that we computed back into the whole matrix
               for (int t = 0; t < numThreads; t++)
					{
						for (int i = 0; i < BLOCK_SIZE; ++i)
							for (int j = 0; j < BLOCK_SIZE; ++j)
								for (int k = 0; k < BLOCK_SIZE; ++k)
								   matrixC_InPPE[(spu_indices[t].i_block + i) * N + (spu_indices[t].j_block + j)] += cb[t].SPU_result[i * BLOCK_SIZE + j];
					}

#ifdef TIMING
	            end_t = __mftb();
	            total_t += (float)((end_t - start_t) / 79.8);
#endif
   
					// Destroy the SPE contexts
					for (int t = 0; t < numThreads; t++)
					{
						if ((rc[t] = spe_context_destroy(data[t].speid)) != 0)
						{
							fprintf(stderr, "Failed spe_context_destroy(rc=%d, errno=%d strerror=%s)\n", rc[t], errno, strerror(errno));
							exit(1);
						}
					}

               thread = 0;
               remainingThreads -= numThreads;
               numThreads = (remainingThreads >= NUM_SPUS) ? NUM_SPUS : remainingThreads;
            }
         }
      }
   }

   // Divide all cells by the block size
   for (int i = 0; i < N; ++i)
	{
		for (int j = 0; j < N; ++j)
		{
         matrixC_InPPE[i * N + j] /= BLOCK_SIZE;
      }
   }

#ifdef TIMING
   printf("\n\n\nTotal time used by SPE: %f us\n\n\n", total_t);
#endif

   // Print the result of multiplying AxB
   //printMatrix("C", matrixC_InPPE);

   // Free up dynamic memory
   free(matrixA_InPPE);
   free(matrixB_InPPE);
   free(matrixC_InPPE);

	return 0;
}

//
// Create and run an SPE thread
//
// @param   void* arg   Function pointer for spe thread
//
void *ppu_pthread_function(void *arg)
{
	ppu_pthread_data_t *datap = (ppu_pthread_data_t *)arg;
	int rc;
	unsigned int entry = SPE_DEFAULT_ENTRY;
	if ((rc = spe_context_run(datap->speid, &entry, 0, datap->argp, NULL, NULL)) < 0)
	{
		fprintf(stderr, "Failed spe_context_run(rc=%d, errno=%d, strerror=%s)\n", rc, errno, strerror(errno));
		exit(1);
	}
	pthread_exit(NULL);
}

// 
// Initialize matrices A and B with random integers from (0, 100]
// and matrix C to 0.
//
void initMatrices()
{
   // Allocate space
   matrixA_InPPE = (int*)malloc(MATRIX_SIZE * sizeof(int));
   matrixB_InPPE = (int*)malloc(MATRIX_SIZE * sizeof(int));
   matrixC_InPPE = (int*)malloc(MATRIX_SIZE * sizeof(int));

	for(int row = 0; row < N; row++)
	{
		for(int col = 0; col < N; col++)
		{
			matrixA_InPPE[row * N + col] = random() % 100;  // Random between (0, 100]
			matrixB_InPPE[row * N + col] = random() % 100;  // Random between (0, 100]
         matrixC_InPPE[row * N + col] = 0;
		}
	}
}

//
// Prints a matrix
//
// @param   char* name     Matrix name
//          int*  matrix   Matrix data
//
void printMatrix(char *name, int *matrix)
{
	printf("Matrix %s:\n", name);
	for(int row = 0; row < N; row++)
	{
		for(int col = 0; col < N; col++)
		{
			printf("%8d", matrix[row * N + col]);
		}
		printf("\n");
	}
}